#!/usr/bin/ruby
# -*- coding: utf-8 -*-
# -------------------------------------------------------------------------------
#    DBGP-Tunnel server implementaion in Ruby
#    http://code.google.com/p/dbgp-tunnel/
# 
#    Copyright (c) 2011 HANAI Tohru <hanai.tohru@gmail.com>
# 
#    Contributor(s): none yet.
# -------------------------------------------------------------------------------
#   This software is provided 'as-is', without any express or implied
#   warranty. In no event will the authors be held liable for any damages
#   arising from the use of this software.
# 
#   Permission is granted to anyone to use this software for any purpose,
#   including commercial applications, and to alter it and redistribute it
#   freely, subject to the following restrictions:
# 
#   1. The origin of this software must not be misrepresented; you must not
#      claim that you wrote the original software. If you use this software
#      in a product, an acknowledgment in the product documentation would be
#      appreciated but is not required.
#   2. Altered source versions must be plainly marked as such, and must not be
#      misrepresented as being the original software.
#   3. This notice may not be removed or altered from any source distribution.
# -------------------------------------------------------------------------------

require 'optparse'
require 'socket'
require 'logger'
require 'monitor'

Version = '0.3'
$logger = Logger.new(STDOUT)
$logger.level = Logger::INFO
STDOUT.sync = true

# commmands/responses
REQUEST_TUNNEL        = "tunnel"
REQUEST_SESSION       = "session"
REQUEST_START         = "start"
RESPONSE_ACCEPT       = "accept"
RESPONSE_OK           = "ok"
RESPONSE_ERROR        = "error"

# errors
ERROR_INVALID_COMMAND  = "invalid_command"
ERROR_SYNTAX           = "syntax_error"
ERROR_IDEKEY_IN_USE    = "idekey_in_use"
ERROR_INVALID_IDEKEY   = "invalid_idekey"
ERROR_UNKNOWN_SESSION  = "unknown_session"

class TCPSocket
  def send(request, param = nil)
    if param.nil?
      write("#{request}\0")
    else
      write("#{request} #{param}\0")
    end
  end

  def recv()
    gets("\0").chomp("\0").split(' ', 2)
  end
end

class DbgpTunnelPipe
  attr_accessor :name, :tunnel_socket, :dbgp_socket, :init_message

  def initialize(logger)
    @logger = logger
    @isRunning = false
  end

  def isRunning?()
    @isRunning
  end

  def close()
    begin
      @dbgp_socket.close if @dbgp_socket and not dbgp_socket.closed?
    rescue
    end

    begin
      @tunnel_socket.close if @tunnel_socket and not @tunnel_socket.closed?
    rescue
    end

    @on_close.call(self) if @on_close
  end

  def run()
    @logger.debug('DbgpTunnelPipe.run()')
    Thread.start do
      begin
        runSession()
      rescue Exception => e
        @logger.debug("tunnel exception: #{$!}")
        puts e.backtrace
      end
      @isRunning = false
      close
    end
  end

  def onClose()
    @on_close = Proc.new
  end

  def runSession()
    if @isRunning
      @logger.warn("DbgpTunnelPipe.run: bug? '#{name}' is already running")
      return
    end
    @isRunning = true

    unless @tunnel_socket
      @logger.warn("DbgpTunnelPipe.run: bug? tunnel_socket is not set")
      return
    end
    unless @dbgp_socket
      @logger.warn("DbgpTunnelPipe.run: bug? dbgp_socket is not set")
      return
    end

    def @tunnel_socket.kind
      'tunnel'
    end

    def @dbgp_socket.kind
      'dbgp'
    end

    @dbgp_socket.binmode
    @tunnel_socket.binmode
    @tunnel_socket.write(@init_message)

    socks = [@tunnel_socket, @dbgp_socket]

    while true
      nsock = IO.select(socks)
      next if nsock.nil?

      for s in nsock[0]
        opposite = (s == @tunnel_socket ? @dbgp_socket : @tunnel_socket)
        if s.eof?
          @logger.debug("#{s.kind} session has disconnected")
          s.close
          opposite.close
          return
        end

        if s == @dbgp_socket
          message = s.gets("\0") + s.gets("\0")
        else
          message = s.gets("\0")
        end
          
        if message.nil?
          @logger.debug("#{s.kind} session has disconnected")
          opposite.close
          return
        end

        @logger.debug("(#{s.kind}) " + message)
        begin
          opposite.write(message)
        rescue
          @logger.debug("#{opposite.kind} session has disconnected")
          s.close
          return
        end
      end
    end
  end
  private :runSession
end

class DbgpTunnelServer
  def initialize(tunnel_addr, tunnel_port,
                 dbgp_addr, dbgp_port)
    @logger = $logger

    @tunnel_listener = TCPServer.open(tunnel_addr, tunnel_port)
    @logger.info("tunnel listener is on #{getSocketName(@tunnel_listener)}")

    @dbgp_listener = TCPServer.open(dbgp_addr, dbgp_port);
    @logger.info("dbgp listener is on #{getSocketName(@dbgp_listener)}")

    @tunnel_catalog = {}
    @tunnel_catalog.extend(MonitorMixin)

    @pipe_catalog = {}
    @pipe_catalog.extend(MonitorMixin)
  end

  def getSocketName(socket)
    addr = socket.addr
    addr.shift
    addr.join(':')
  end

  def run()
    socks = [@tunnel_listener, @dbgp_listener]
    while true
      nsock = IO.select(socks)
      next if nsock.nil?

      for s in nsock[0]
        socket = s.accept

        case s
        when @tunnel_listener
          @logger.debug("tunnel listener has accepted a new connection")
          startTunnelSession(socket)
        when @dbgp_listener
          @logger.debug("dbgp listener has accepted a new connection")
          startDbgpSession(socket)
        end
      end
    end
  end

  def startTunnelSession(socket)
    socket.binmode

    Thread.start do
      begin
        startTunnelSession1(socket)
      rescue Exception => e
        @logger.debug("tunnel exception: #{$!}")
        puts e.backtrace
      end
    end
  end
  
  def startTunnelSession1(socket)
    idekey = nil

    while true
      nsock = IO.select([socket])
      next if nsock.nil?

      if socket.eof?
        @logger.info("tunnel (idekey: '#{idekey}') has disconnected")
        @tunnel_catalog.synchronize do
          @tunnel_catalog.delete(idekey)
        end
        return
      end

      @logger.debug('reading request from tunnel')
      (request, param) = socket.recv
      @logger.debug("received: '#{request} #{param}'")

      case request
      when REQUEST_TUNNEL
        param.strip! if not param.nil?
        if param.nil? or param == ''
          @logger.info("tunnel sends an invalid idekey '#{param}'")
          socket.send(RESPONSE_ERROR, ERROR_INVALID_IDEKEY)
          socket.close
          return
        elsif @tunnel_catalog.key?(param)
          @logger.info("tunnel (idekey: '#{param}') is in use; new connection is rejected")
          socket.send(RESPONSE_ERROR, ERROR_IDEKEY_IN_USE)
          socket.close
          return
        else
          idekey = param
          @logger.debug("tunnel (idekey: '#{idekey}') has connected")
          @tunnel_catalog.synchronize do
            @tunnel_catalog[idekey] = { 'tunnel' => socket }
          end
          socket.send(RESPONSE_ACCEPT, idekey)
        end
      when REQUEST_START
        if @pipe_catalog.key?(param)
          pipe = @pipe_catalog[param]
          if not pipe.isRunning?
            pipe.tunnel_socket = socket
            socket.send(RESPONSE_ACCEPT, param)
            pipe.run
            return
          end
        end

        @logger.info("tunnel sends an invalid session key #{param}")
        socket.send(RESPONSE_ERROR, ERROR_UNKNOWN_SESSION)
        sleep(1)
        socket.close
        return
      else
        @logger.info("tunnel sends unknown data: '#{request} #{param}'")
      end
    end
  end

  def parseIdeKey(message)
    return $1 if message.match(/idekey="(.*?)"/)
  end

  def findTunnelSocketByIdeKey(idekey)
    @tunnel_catalog.synchronize do
      if @tunnel_catalog.key?(idekey)
        return @tunnel_catalog[idekey]['tunnel']
      end
    end
  end

  def startDbgpSession(dbgp_socket)
    if not requestDbgpBinding(dbgp_socket)
      dbgp_socket.close
    end
  end

  def requestDbgpBinding(dbgp_socket)
    (peer_port, peer_addr) = Socket.unpack_sockaddr_in(dbgp_socket.getpeername)

    while true
      nsock = IO.select([dbgp_socket])
      next if nsock.nil?

      for s in nsock[0]
        if s.eof?
          @logger.info("new dbgp session has immidiately disconnected")
          return false
        end

        len = s.gets("\0")
        message = s.gets("\0") if not len.nil?

        if message.nil?
          @logger.info("new dbgp session has immidiately disconnected")
          return false
        end

        idekey = parseIdeKey(message)
        if idekey.nil?
          @logger.info("failed to parse an idekey from a message: #{message}")
          return false
        end
        @logger.debug("new debgp's idekey: '#{idekey}'")

        tunnel_socket = findTunnelSocketByIdeKey(idekey)
        if not tunnel_socket
          @logger.info("no tunnel for the dbgp session")
          return false
        end
          
        pipe = DbgpTunnelPipe.new(@logger)
        pipe.name = s.object_id.to_s
        pipe.dbgp_socket = s
        pipe.init_message = "#{len}#{message}"

        pipe.onClose do
          @pipe_catalog.synchronize do
            @pipe_catalog.delete(pipe.name)
          end
        end

        @pipe_catalog[pipe.name] = pipe

        begin
          tunnel_socket.send(REQUEST_SESSION, pipe.name)
          return true
        rescue
          @logger.info("failed to send a session connection request to a tunnel")
          pipe.close
          return false;
        end
      end
    end
  end
end

def main
  tunnel_addr = 'localhost'
  tunnel_port = 9100
  dbgp_addr = 'localhost'
  dbgp_port = 9000

  opt = OptionParser.new
  opt.on('-t', '--tunnel=ADDR:PORT', 'interface address and port for a tunnel') do |v|
    (tunnel_addr, tunnel_port) = v.split(':')
  end

  opt.on('-s', '--session=ADDR:PORT', 'interface address and port for a DBGP session') do |v|
    (dbgp_addr, dbgp_port) = v.split(':')
  end

  opt.on('-d', '--debug', 'output log for debug') do |v|
    $logger.level = Logger::DEBUG
  end

  opt.parse!(ARGV)
  
  tunnel_server = DbgpTunnelServer.new(tunnel_addr, tunnel_port, dbgp_addr, dbgp_port)
  tunnel_server.run
end

main()
